from flask import Flask, request, jsonify
import os
#import speech_recognition as sr
from funasr import AutoModel

from funasr.utils.postprocess_utils import rich_transcription_postprocess


from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

from flask_cors import *

app = Flask(__name__)
CORS(app, resources=r'/*')



# ======================= 语音转文字，中文语音 =======================
model_zh = AutoModel(model="./used_weights/modelscope/hub/models/iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch", model_revision="v2.0.4",
                  vad_model="./used_weights/modelscope/hub/models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch", vad_model_revision="v2.0.4",
                  punc_model="./used_weights/modelscope/hub/models/iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch", punc_model_revision="v2.0.4",
                  disable_update=True,
                  cache_dir="./used_weights",
                  force_download=False,
                  resume_download=False,
                  # spk_model="cam++", spk_model_revision="v2.0.2",
                  )

# ======================= 语音转文字，英文语音 =======================


#model_dir = "iic/SenseVoiceSmall"
model_dir = "./used_weights/modelscope/hub/iic/SenseVoiceSmall"

model_en = AutoModel(
    model=model_dir,
    trust_remote_code=True,
    remote_code="./model.py",
    disable_update=True,
    vad_model="./used_weights/modelscope/hub/models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
    vad_kwargs={"max_single_segment_time": 30000},
    device="cuda:0",
)



# ======================= 翻译，中文转英文 =======================
pipeline_trans_zh2en = pipeline(task=Tasks.translation, model="./used_weights/modelscope/hub/models/damo/nlp_csanmt_translation_zh2en")


# ======================= 翻译，英文转中文 =======================
pipeline_trans_en2zh = pipeline(task=Tasks.translation, model="./used_weights/modelscope/hub/models/damo/nlp_csanmt_translation_en2zh")





def transcribe_audio_zh(file_path):
    res = model_zh.generate(input=file_path,
                     batch_size_s=300)
    return res



def transcribe_audio_en(file_path):
    res = model_en.generate(
        input=file_path,
        cache={},
        language="en",  # "zn", "en", "yue", "ja", "ko", "nospeech"
        use_itn=True,
        batch_size_s=60,
        merge_vad=True,  #
        merge_length_s=15,
        )
    #text = rich_transcription_postprocess(res[0]["text"])

    return res
    
    # text = model_en(file_path)
    # return text






#app = Flask(__name__)

# 配置上传文件的保存路径
UPLOAD_FOLDER = 'uploads'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

@app.route('/speech2text_zh', methods=['POST'])
def speech2text_zh():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    
    file = request.files['file']
    
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400
    
    if file and file.filename.endswith('.mp3'):
        # 保存文件到本地
        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
        file.save(file_path)
        
        # 调用语音识别函数
        try:
            text = transcribe_audio_zh(file_path)
            return jsonify({"text": text}), 200
        except Exception as e:
            return jsonify({"error": f"Transcription failed: {str(e)}"}), 500
        finally:
            # 删除临时文件
            if os.path.exists(file_path):
                os.remove(file_path)
    else:
        return jsonify({"error": "Invalid file format. Only MP3 files are supported."}), 400


@app.route('/speech2text_en', methods=['POST'])
def speech2text_en():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    
    file = request.files['file']
    
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400
    
    if file and file.filename.endswith('.mp3'):
        # 保存文件到本地
        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
        file.save(file_path)
        
        # 调用语音识别函数
        try:
            text = transcribe_audio_en(file_path)
            return jsonify({"text": text}), 200
        except Exception as e:
            return jsonify({"error": f"Transcription failed: {str(e)}"}), 500
        finally:
            # 删除临时文件
            if os.path.exists(file_path):
                os.remove(file_path)
    else:
        return jsonify({"error": "Invalid file format. Only MP3 files are supported."}), 400




@app.route('/translate', methods=['POST'])
def translate():
    try:
        # 获取 POST 请求中的 JSON 数据
        data = request.get_json()
        if not data or 'text' not in data or 'dest_lang' not in data:
            return jsonify({"error": "Missing 'text' or 'dest_lang' in request"}), 400
        
        text_to_translate = data['text']
        dest_language = data['dest_lang']
        
        print("text_to_translate: ",text_to_translate)
        if text_to_translate.strip() == "":
            return jsonify({
                "original_text": "",
                "translated_text": "",
                "destination_language": dest_language
            }), 200

        if dest_language == "zh":
            # 调用翻译 API
            outputs = pipeline_trans_en2zh(input=text_to_translate)
            translated = outputs['translation']
            
            # 返回翻译结果
            return jsonify({
                "original_text": text_to_translate,
                "translated_text": translated,
                "destination_language": dest_language
            }), 200
        elif dest_language == "en":
            # 调用翻译 API
            outputs = pipeline_trans_zh2en(input=text_to_translate)
            translated = outputs['translation']
            
            # 返回翻译结果
            return jsonify({
                "original_text": text_to_translate,
                "translated_text": translated,
                "destination_language": dest_language
            }), 200
            

    except Exception as e:
        return jsonify({"error": str(e)}), 500



if __name__ == '__main__':
    app.run(debug=False, host="0.0.0.0",port=5072)
